//
// main.rs
// Copyright (C) 2020 zhuzhengyi1 <zhuzhengyi@jd.com>
// Distributed under terms of the MIT license.
//

//
use async_std::task;
use std::collections::HashMap;
use std::fs::File;
use std::io::prelude::*;
use std::io::BufReader;

use clap::Clap;

mod ssh_node;

#[derive(Clap)]
#[clap(name = "pd")]
struct Opts {
    #[clap(short, long)]
    quiet: bool,

    #[clap(short = "N", long)]
    node_file: Option<String>,

    #[clap(short = "n", long)]
    nodes: Option<String>,

    #[clap(short, long)]
    passwd: Option<String>,

    #[clap(short, long, default_value = "root")]
    user: String,

    #[clap(subcommand)]
    subcmd: SubCommand,
}

#[derive(Clap)]
enum SubCommand {
    Run(Run),
    Upload(Upload),
}

#[derive(Clap)]
struct Run {
    #[clap(short, long, required = true)]
    cmd: String,
}

#[derive(Clap)]
struct Upload {
    #[clap(short, long, required = true)]
    src_file: String,

    #[clap(short, long, required = true)]
    tgt_file: String,
}

//
fn main() {
    let mut ssh_nodes: HashMap<String, ssh_node::SSHNode> = HashMap::new();
    let mut hosts: HashMap<String, String> = HashMap::new();

    let opts: Opts = Opts::parse();

    if let Some(node_file) = &opts.node_file {
        let f = File::open(&node_file).unwrap();
        let buffered = BufReader::new(f);
        for host in buffered
            .lines()
            .map(|l| l.unwrap())
            .filter(|l| !(l.starts_with("#") || l.is_empty()))
        {
            if !hosts.contains_key(&host) {
                hosts.insert(host.clone(), host.clone());
            }
        }
    }

    if let Some(nodes) = &opts.nodes {
        for host in nodes.split(",").map(|h| h.to_string()) {
            if !hosts.contains_key(&host) {
                hosts.insert(host.clone(), host.clone());
            }
        }
    }

    //for (_, node) in ssh_nodes.iter_mut() {
    //    node.init(&opts.user, opts.passwd.clone()).unwrap();
    //}
    let mut tasks = vec![];
    for (_, host) in hosts.into_iter() {
        let user = opts.user.clone();
        let passwd = opts.passwd.clone();
        tasks.push(task::spawn(async move {
            (host.clone(), ssh_node::SSHNode::new(&host, &user, passwd))
        }));
    }

    task::block_on(async {
        for t in tasks {
            let (host, node) = t.await;
            if !ssh_nodes.contains_key(&host) {
                ssh_nodes.insert(host, node);
            }
        }
    });

    match opts.subcmd {
        SubCommand::Run(run) => {
            if !&opts.quiet {
                println!("run: {}", &run.cmd);
            }

            let mut tasks = vec![];
            for (_, node) in ssh_nodes.into_iter() {
                let cmd = String::from(&run.cmd);
                tasks.push(task::spawn(
                    async move { (node.host.clone(), node.exec(&cmd)) },
                ));
            }
            task::block_on(async {
                for t in tasks {
                    let (host, r) = t.await;
                    match r {
                        Ok(r) => {
                            println!("{}: \n{}", &host, r);
                        }
                        Err(err) => println!("{} error: {}", &host, err),
                    }
                }
            });
        }
        SubCommand::Upload(upload) => {
            if !&opts.quiet {
                println!("upload file");
            }
            let mut tasks = vec![];
            for (_, node) in ssh_nodes.into_iter() {
                let src = String::from(&upload.src_file);
                let tgt = String::from(&upload.tgt_file);
                tasks.push(task::spawn(async move {
                    (node.host.clone(), node.upload(&src, &tgt))
                }));
            }
            task::block_on(async {
                for t in tasks {
                    let (host, r) = t.await;
                    match r {
                        Ok(_) => {
                            println!("{} upload success", &host);
                        }
                        Err(err) => println!("{} upload failed: {}", &host, err),
                    }
                }
            });
        }
    }
}
